% Experiment Setup 2: Classification using training data samples

addpath(genpath('./utility/'))
global numBins

classLabels = ['A' 'B' 'C' 'D' 'F' 'G' 'H' 'K' 'M' 'O' 'P' 'S' 'Z']; % note: currently do not use W - water bottle
C = numel(classLabels);
trainPath = './images/training/';
masksPath = './images/trainingmasks/';

testImgsPath = './images/test/easy/';
colorSpace = 'hsi';
classifier = 'svm';

numBins = 100;

list = subdir([trainPath '*.jpg']); list = {list(:).name};
N = numel(list);
% Build the training models
if ~exist(['./data/trainingSamples_' colorSpace '.mat'])
    tic;
    trainX = -ones(N,numBins*5); % training samples 
    trainY = -ones(N,1); % training labels
    % extract features for each training image
    for i=1:N
        imgName = list{i};
        I = imread(imgName);
        mask = load([masksPath imgName(end-7:end-4)]);
        [chist1,chist2,chist3] = Color_Features(I,mask.mask,colorSpace,true);
        [thist1,thist2] = Texture_Features(mask.mask,I,1);
        trainX(i,:) = horzcat(chist1,chist2,chist3,thist1,thist2);
        trainY(i) = strfind(classLabels, imgName(end-7));
        disp(['Extracted features for training sample ' num2str(i) ' of ' num2str(N) '.'])
    end
    time_extractTrain = toc;
    save(['./data/trainingSamples_' colorSpace], 'trainX', 'trainY', 'time_extractTrain')
else
    load(['./data/trainingSamples_' colorSpace])
end

% Build the training mean HS models to be used for segmentation
if ~exist('./data/models_meansHS.mat')
    models_meansHS = -ones(C,2);
    for i=1:C
        classPath = [trainPath classLabels(i) '/'];
        list = dir([classPath '*.jpg']);
        list = {list(:).name};
        
        sumPixels = 0;
        sumHue = 0;
        sumSat = 0;
        for j=1:numel(list)
            imgName = list{j};
            I = imread([classPath imgName]);
            mask = load([masksPath imgName(1:end-4)]); mask = mask.mask;
            hsv = rgb2hsv(I);
            H = hsv(:,:,1);
            S = hsv(:,:,2);
            sumPixels = sumPixels + sum(mask(:));
            sumHue = sumHue + sum(H(mask));
            sumSat = sumSat + sum(S(mask));
        end
        models_meansHS(i,1) = sumHue/sumPixels;
        models_meansHS(i,2) = sumSat/sumPixels;
        disp(['Mean hue and saturation saved for class ' classLabels(i) '.'])
    end
    save('./data/models_meansHS.mat', 'models_meansHS');
else
    load('./data/models_meansHS.mat')
end

% Apply PCA to reduce feature dimensionality and retain 90% variance
[p,c,d,trainXPCA] = lowDimPCA(trainX,0.99);
% Apply MDA to further reduce dimensionality
pm = mda(trainXPCA, trainY, C-1);
trainXMDA = (trainXPCA * pm);


% Make a mask of true classes for the test images
list = dir([testImgsPath '*.jpg']); list = {list(:).name};
T = numel(list); % T = number of test images
testClasses = zeros(1,T);
for i=1:numel(list)
    parts = textscan(list{i}, '%s%s%s', 'Delimiter', '-.');
    testClasses(i) = strfind(classLabels, char(parts{1}));
end
truthMask = zeros(C,T);
truthMask(sub2ind(size(truthMask), testClasses, 1:T)) = 1;


testEvalMat = load('./data/testEvalMat_easy'); % note: this is used to help with ground truth for calculating accuracies
% Classify each test image
scores = zeros(C,T);
tic;

for t=1:T
    testFile = [testImgsPath list{t}];
    I = imread(testFile);
    % Get mask(s) from segmentation results
    ROIs = segmentationL(I, models_meansHS);
    
%     subplot(1,size(ROIs,3)+1,1)
%     imshow(I)
    
    testX = zeros(T,numBins*5);
    testF = zeros(1,numBins*5);
    for r=1:size(ROIs,3)
        mask = ROIs(:,:,r);
        
%         subplot(1,size(ROIs,3)+1,r+1)
%         imshow(mask)
        
        [chist1,chist2,chist3] = Color_Features(I,mask,colorSpace,true);
        [thist1,thist2] = Texture_Features(mask,I,1);
        x = horzcat(chist1,chist2,chist3,thist1,thist2);
        
        trueLoc = testEvalMat.locations(strmatch(list{t},testEvalMat.filenames,'exact'),:);
        trueLoc = round(trueLoc);
        if mask(sub2ind(size(mask),trueLoc(1),trueLoc(2)))==1
            testX(t,:) = x;
        else 
            testF = vertcat(testF,x);
        end
    end
    disp(['Extracted features for test image ' num2str(t) ' of ' num2str(T) '.'])
end

testXPCA = (testX - repmat(c, size(testX,1), 1)) * p(:,1:d);
testXMDA = testXPCA * pm;
testFPCA = (testF - repmat(c, size(testF,1), 1)) * p(:,1:d);
testFMDA = testFPCA * pm;

scoresX = ClassifyTestSamples(trainXMDA, trainY, testXMDA, testClasses, classifier);
scoresF = ClassifyTestSamples(trainXMDA, trainY, testFMDA, testClasses, classifier);

time_test = toc;

save(['./results/' colorSpace '_' classifier], 'scoresX', 'scoresF', 'truthMask')

% Calculate accuracy
[~, classIds] = sort(scores,'descend');
rank = 1;
accuracy = sum(truthMask(sub2ind(size(truthMask), ids(rank,:), 1:T)) > 0) / T;
correct = truthMask(sub2ind(size(truthMask), ids(rank,:), 1:T)) > 0;


% Find misclassified images


